import os

import cv2
import numpy as np
from loguru import logger
from scipy import optimize


def remove_small_components(img, threshold):
    """
    删除小的连通域
    :param img:
    :param threshold:
    :return:
    """
    # 获取连通域
    _, labels, stats, _ = cv2.connectedComponentsWithStats(img, connectivity=4)
    # 遍历所有连通域
    for i in range(1, np.max(labels) + 1):
        # 计算连通域的面积
        area = stats[i, cv2.CC_STAT_AREA]
        # 去除面积小于某个阈值的连通域
        if area < threshold:
            img[labels == i] = 0
    return img


def get_diameter_hough(src):
    # 转换为灰度图像
    gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    # 应用高斯模糊来降噪, 暂时不用
    # blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    # 二值化，小于200的归为黑色
    ret, image = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    image = remove_small_components(image, 1000)

    # 使用霍夫圆变换来检测圆形
    circles = cv2.HoughCircles(image, cv2.HOUGH_GRADIENT, 1, minDist=1000, param1=100, param2=15, minRadius=980,
                               maxRadius=1000)
    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    # 将检测到的圆形转换为整数
    if circles is not None:
        circles = np.uint16(np.around(circles))
        for x, y, r in circles[0, :]:
            # 在原图像上绘制圆形
            cv2.circle(src, (x, y), r, (0, 255, 0), 2)
            cv2.circle(image, (x, y), r, (0, 255, 0), 2)
            result_str = f"Circle detected at: ({x}, {y}), radius: {r}, diameter: {r * 2}"
            # 输出圆形的参数
    else:
        result_str = "no circle found"

    return result_str, src, image


def get_diameter_leastsq(xdata, ydata):
    """

    :param xdata:
    :param ydata:
    :return:
    """

    # 定义目标函数，计算误差平方和
    def error_func(params, xdata, ydata):
        xc, yc, r = params
        return (xdata - xc) ** 2 + (ydata - yc) ** 2 - r ** 2

    # 假设我们有一组(x, y)坐标点，这些点大致分布在某个圆上
    # xdata = np.array([1, 2, 3, 4, 5])
    # ydata = np.array([1, 3, 1, 3, 1])

    # 初始参数猜测
    x_m = float(np.mean(xdata))
    y_m = float(np.mean(ydata))
    logger.debug(f"x_m:{x_m}")
    center_estimate = np.array([x_m, y_m, 700])
    # 使用leastsq函数进行拟合
    result = optimize.leastsq(error_func, center_estimate, args=(xdata, ydata))
    # 获取拟合结果
    xc, yc, r = result[0]

    print("圆心坐标：(%.2f, %.2f)" % (xc, yc))
    print("半径：%.2f" % r)
    print(f"直径：{r * 2:.2f}")
    return (xc, yc), r


def test_HoughCircles():
    test_folder_path = "../resource/static/warmup/crop_test"
    result_folder_path = "../resource/static/warmup/results"
    for filename in os.listdir(test_folder_path):
        src = cv2.imread(f'{test_folder_path}/{filename}', cv2.IMREAD_COLOR)
        result_str, src, image = get_diameter_hough(src)
        print(f"{filename}: {result_str}")
        cv2.putText(src, result_str, (100, 100), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 255, 0), 1)
        cv2.putText(image, result_str, (100, 100), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 255, 0), 1)
        cv2.imwrite(f'{result_folder_path}/processed-src-{filename}', src)
        cv2.imwrite(f'{result_folder_path}/processed-bi-{filename}', image)


def get_edges_canny(img):
    """
    获取边缘.
    :param img:
    :return:
    """
    min_thresh_canny = 57
    max_thresh_canny = 141
    canny_flag_l2gradient = False
    # img_blurred = cv2.GaussianBlur(img.copy(), (bsize_pos * 2 + 1, bsize_pos * 2 + 1), bamp_pos)
    edges = cv2.Canny(img, min_thresh_canny, max_thresh_canny,
                      L2gradient=canny_flag_l2gradient)
    logger.debug(f"edges:{edges}")
    edge_idx = np.where(edges != 0.)
    logger.debug(f"edge_idx:{edge_idx}")
    return edge_idx, edges


def paint_circle(image, center, radius):
    cv2.circle(image, center, int(radius), (0, 255, 0), 2)
    return image


def preprocess_image(src):
    """
    预处理
    :param src:
    :return:
    """
    gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    # 应用高斯模糊来降噪
    # blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    ret, image = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    return image


if __name__ == '__main__':
    src = cv2.imread(f'../resource/static/warmup/crop_test/crop1.jpg', cv2.IMREAD_COLOR)
    image = preprocess_image(src)
    edge_idx, edges = get_edges_canny(image)
    # image[edges!= 0] = [0, 0, 255]
    print(edges != 0)
    src[edges != 0] = [0, 0, 255]  # 用红色标出edge
    center, r = get_diameter_leastsq(edge_idx[1], edge_idx[0])
    # center, r = regress_circle(edge_idx[1], edge_idx[0])
    xc, yc = center

    paint_circle(src, (int(xc), int(yc)), r)
    cv2.imshow("result", src)
    cv2.waitKey(0)
